import datetime
import os
import types
import pandas as pd

from pymongo.cursor import Cursor

from app.config.settings import BASE_DIR
# from run import app as start_app
from app.db.models.users import User
from app.handler.v2_index import LastTask
from app.libs.enums import TASK_TYPE_TUPLE


class LastTaskContainDomain(LastTask):
    allow_key = ('endTime', 'taskType', 'domain')


class UserTasks:
    """用户最新监测数据导出"""
    def __init__(self, username, task_type):
        self.username = username
        self.task_type = task_type

    @property
    def user(self):
        return User.objects.filter(username=self.username).first() or {}

    @property
    def uid(self):
        if not self.user:
            raise ValueError("用户不存在")
        return self.user.id

    @property
    def filename(self):
        """目前只导出漏洞监测详情"""
        _type = "漏洞"
        now = datetime.datetime.now().strftime("%Y%m%d%H")
        return f"ScanV{_type}监测详情_{self.username}_{now}.csv"

    def export(self, to_path=None) -> Cursor:
        type_data = self._get_type_data()
        export_method = f"export_{self.task_type}"
        if hasattr(self, export_method):
            return getattr(self, export_method)(type_data, to_path)
        else:
            raise AttributeError("该类型处理方法未定义")

    def _get_type_data(self):
        if self.task_type == "all":
            task_types = TASK_TYPE_TUPLE
        else:
            task_types = [self.task_type]

        type_data = LastTaskContainDomain(uid=self.uid, task_type_list=task_types).get_warnings_group_task()
        return type_data

    def _format_vul(self, vul_data) -> types.GeneratorType:
        for task in vul_data:
            warnings = task.get("warnings", [])
            for warning in warnings:
                traffics = warning.get("traffics") or [{}]
                request_body = traffics[0].get("request", {})
                response_body = traffics[0].get("response", {})

                status = "已处理" if warning.get("isResolved") else "已忽略" if warning.get("ignoredTime") else "未处理"
                found_at = warning.get("foundAt").strftime("%Y-%m-%d") if warning.get("foundAt") else None
                export_dict = {
                    "target": task["target"],  # 资产地址
                    "name": task["note"],  # 资产名称
                    "domain": task.get("domain"),  # 域名
                    "levelName": warning.get("levelName"),  # 危险等级名称
                    "level": warning.get("level"),  # 危险等级
                    "title": warning.get("title"),  # 漏洞名称
                    "foundAt": found_at,  # 发现时间
                    "category": warning.get("category"),  # 漏洞风险类别
                    "description": warning.get("description"),  # 漏洞描述
                    "detail": warning.get("detailText"),  # 漏洞详情
                    "impact": warning.get("impact"),  # 漏洞影响
                    "recommendation": warning.get("recommendation"),  # 处理建议
                    "vul_id": str(warning.get("_id")),  # 漏洞id
                    "reference": "\n".join(warning.get("reference", [])),  # 参考信息
                    "request": request_body or None,  # 请求体
                    "response": response_body or None,  # 响应体
                    "payload": warning.get("payload"),  # 载体
                    "status": status  # 处理状态
                }
                yield export_dict

    def export_vul(self, vul_data, to_path=None):
        if not to_path:
            export_dir = os.path.join(BASE_DIR, "export/")
            if not os.path.exists(export_dir):
                os.mkdir(export_dir)
            to_path = os.path.join(export_dir, self.filename)
        # 列顺序及中文映射
        m_headers = [
            ("target", "风险URL地址"),
            ("name", "资产名称"),
            ("domain", "域名"),
            ("vul_id", "漏洞id"),
            ("title", "漏洞名称"),
            ("category", "风险类型"),
            ("level", "漏洞等级"),
            ("levelName", "漏洞等级名"),
            ("foundAt", "发现时间"),
            ("status", "处理状态"),
            ("description", "漏洞描述"),
            ("detail", "漏洞详情"),
            ("impact", "漏洞影响"),
            ("recommendation", "修复建议"),
            ("reference", "参考信息"),
            ("request", "请求体"),
            ("response", "响应体"),
            ("payload", "载荷")
        ]
        vul_lines = self._format_vul(vul_data)
        columns_sequence = []
        rename_mapping = {}
        for en_name, cn_name in m_headers:
            columns_sequence.append(cn_name)
            rename_mapping[en_name] = cn_name
        df = pd.DataFrame(vul_lines)
        df.rename(rename_mapping, axis="columns", inplace=True)
        df.to_csv(to_path, index=False, columns=columns_sequence)
        return to_path
